# query.py
import os
from dotenv import load_dotenv
from pymilvus import connections, Collection
import dashscope  # 直接使用dashscope库
from openai import OpenAI  # 仅用于DeepSeek

# 加载环境变量
load_dotenv()

# 初始化dashscope
dashscope.api_key = os.getenv('DASHSCOPE_API_KEY')

# 初始化DeepSeek客户端
deepseek_client = OpenAI(
    api_key=os.getenv('DEEPSEEK_API_KEY'),
    base_url="https://api.deepseek.com/v1"
    # 注意：OpenAI客户端不支持proxies参数
)

# 配置Milvus连接
connections.connect(
    host=os.getenv('MILVUS_HOST', 'localhost'),
    port=os.getenv('MILVUS_PORT', '19530')
)

# 集合名称
collection_name = "bge_docs"
collection = Collection(collection_name)

def get_embedding(text: str, model="text-embedding-v2"):
    """调用通义千问的Embedding API生成向量"""
    try:
        response = dashscope.TextEmbedding.call(
            model=model,
            input=text
        )
        if response.status_code == 200 and response.output is not None:
            return response.output['embeddings'][0]['embedding']
        else:
            print(f"生成Embedding失败: {response.message}")
            return None
    except Exception as e:
        print(f"生成Embedding时出错: {e}")
        return None

def search_similar_chunks(query: str, top_k: int = 3):
    """在Milvus中搜索相似的文本块，并进行去重"""
    query_vector = get_embedding(query)
    if not query_vector:
        return []
    
    # 搜索参数
    search_params = {"metric_type": "L2", "params": {"ef": 64}}
    
    # 执行搜索 - 获取更多结果以便去重后仍有足够数量
    results = collection.search(
        data=[query_vector],
        anns_field="vector",
        param=search_params,
        limit=top_k * 2,  # 获取双倍数量的结果
        output_fields=["file_name", "text_chunk"]
    )
    
    retrieved_chunks = []
    seen_texts = set()  # 用于去重的文本集合
    
    for hits in results:
        for hit in hits:
            text_chunk = hit.entity.get('text_chunk')
            # 简单去重：基于文本块内容
            if text_chunk not in seen_texts:
                seen_texts.add(text_chunk)
                chunk_info = {
                    "file_name": hit.entity.get('file_name'),
                    "text_chunk": text_chunk,
                    "distance": hit.distance,
                    "score": 1 - hit.distance  # 转换为相似度分数
                }
                retrieved_chunks.append(chunk_info)
                # 如果已经收集了足够的去重结果，可以提前退出
                if len(retrieved_chunks) >= top_k:
                    break
        if len(retrieved_chunks) >= top_k:
            break
    
    # 按相似度分数排序
    retrieved_chunks.sort(key=lambda x: x['score'], reverse=True)
    # 确保只返回top_k个结果
    return retrieved_chunks[:top_k]

def generate_answer(question: str, context_chunks: list):
    """调用DeepSeek生成答案"""
    if not context_chunks:
        return "抱歉，没有找到相关的信息来回答这个问题。"
    
    # 构建上下文字符串
    context_str = "\n\n".join([
        f"【文档: {chunk['file_name']}, 相关度: {chunk['score']:.3f}】\n{chunk['text_chunk']}" 
        for chunk in context_chunks
    ])
    
    # 构建Prompt
    prompt = f"""请你作为一个专业的文档分析助手，根据以下提供的背景资料，准确、专业地回答用户的问题。

【背景资料】
{context_str}

【用户问题】
{question}

【回答要求】
1. 严格根据背景资料的内容进行回答，不要编造资料中没有的信息。
2. 如果资料中没有答案，请直接说明“根据现有资料无法回答该问题”。
3. 回答要自然、专业、简洁明了。
4. 在回答末尾注明答案的主要来源文档。"""

    try:
        response = deepseek_client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {"role": "system", "content": "你是一个专业、准确的文档分析助手。"},
                {"role": "user", "content": prompt}
            ],
            temperature=0.1,
            max_tokens=2000
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"调用DeepSeek API时出错: {e}"

def main():
    """主函数：交互式问答"""
    # 确保集合已加载
    collection.load()
    print("=" * 60)
    print("🤖 企业知识库问答系统已启动")
    print("=" * 60)
    print(f"集合统计: {collection.num_entities} 条数据已就绪")
    print("输入 'quit' 或 'exit' 退出程序\n")
    
    while True:
        try:
            user_question = input("🙋 请输入您的问题: ").strip()
            
            if user_question.lower() in ['quit', 'exit', '退出']:
                print("再见！")
                break
                
            if not user_question:
                continue
            
            print("🔍 正在检索相关信息...")
            # 搜索相关文本块
            similar_chunks = search_similar_chunks(user_question, top_k=3)
            
            print("🧠 正在生成答案...")
            # 生成答案
            answer = generate_answer(user_question, similar_chunks)
            
            # 显示结果
            print("\n" + "=" * 60)
            print("💡 答案：")
            print(answer)
            print("=" * 60)
            
            # 显示参考来源
            if similar_chunks:
                print("\n📚 参考来源：")
                for i, chunk in enumerate(similar_chunks, 1):
                    print(f"{i}. 《{chunk['file_name']}》 (相关度: {chunk['score']:.3f})")
            print()  # 空行
            
        except KeyboardInterrupt:
            print("\n\n程序已中断。")
            break
        except Exception as e:
            print(f"发生错误: {e}")

if __name__ == "__main__":
    main()